#include "graph.h"
#include "vkinfer/core/logger.h"
#include "vkinfer/parser/onnxparser.h"
#include "vkinfer/runtime/vk/vkcontext.h"
#include "vkinfer/core/topology_sort.h"

namespace vkinfer
{
    Graph::Graph(const std::string& model_file)
    {
        is_valid = false;

        // initialize context
        context = std::make_shared<VkContext>();

        // load model
        OnnxParser parser;
        bool success = parser.open(model_file);
        if (!success)
        {
            VK_ERROR("failed to load onnx file: %s\n", model_file.c_str());
        }
        // load io info
        parser.load_io_info(input_names, output_names);

        // load weights
        std::map<std::string, std::shared_ptr<OnnxWeight>> weights;
        parser.load_weights(weights);
        
        // load op info and create layers
        std::map<std::string, std::shared_ptr<OnnxNodeInfo>> node_infos;
        parser.load_network(node_infos);

        // sort layers
        sorted_layers = topology_sort(node_infos, input_names);

        // create blobs
        for (auto& info : node_infos)
        {
            for (auto& name : info.second->input_names)
            {
                if (blobs.find(name) != blobs.end())
                    continue;
                BlobTransferDirection direction = BlobTransferDirection::BLOB_TRANSFER_NONE;
                if (std::find(input_names.begin(), input_names.end(), name) != input_names.end())
                    direction = BlobTransferDirection::BLOB_TRANSFER_DST;
                auto blob = context->make_blob(name, direction);
                blobs.insert(std::make_pair(name, blob));
            }

            for (auto& name : info.second->output_names)
            {
                if (blobs.find(name) != blobs.end())
                    continue;
                BlobTransferDirection direction = BlobTransferDirection::BLOB_TRANSFER_NONE;
                if (std::find(output_names.begin(), output_names.end(), name) != output_names.end())
                    direction = BlobTransferDirection::BLOB_TRANSFER_SRC;
                auto blob = context->make_blob(name, direction);
                blobs.insert(std::make_pair(name, blob));
            }
        }

        // create weight_blobs
        std::map<std::string, std::shared_ptr<Blob>> weight_blobs;
        for (auto& wts : weights)
        {
            std::string name = wts.first;
            std::shared_ptr<OnnxWeight> weight = wts.second;
            auto blob = context->make_blob(name, BlobTransferDirection::BLOB_TRANSFER_DST);
            blob->copy_from((float*)weight->values, weight->shape);
            weight_blobs.insert(std::make_pair(name, blob));
        }

        // create layers
        for (auto& pinfo : node_infos)
        {
            std::string name = pinfo.first;
            std::shared_ptr<OnnxNodeInfo> info = pinfo.second;
            std::shared_ptr<Operator> op = context->make_op(name, info);
            // set weights
            if (!op)
            {
                VK_WARN("failed to create op: %s\n", name.c_str());
                return;
            }
            // check if the layer has no deps
            bool is_head = true;
            for (auto& inp_name : info->input_names)
            {
                if (std::find(input_names.begin(), input_names.end(), inp_name) == input_names.end())
                {
                    is_head = false;
                    break;
                }
            }
            op->top_layer = is_head;
            
            for (auto& weight_name : info->weight_names)
            {
                if (weight_blobs.find(weight_name) != weight_blobs.end())
                    op->set_weight(weight_name, weight_blobs[weight_name]);
            }
            operators.insert(std::make_pair(name, op));
        }
        // network construction ready.
        is_valid = true;
    }

    bool Graph::feed_input(const std::string& name, const std::vector<float>& data)
    {
        auto it = blobs.find(name);
        if (it == blobs.end())
            return false;

        auto blob = blobs[name];
        blob->copy_from(data);
        return true;
    }

    bool Graph::feed_input(const std::string& name, const std::vector<float>& data, const std::vector<uint32_t>& shape)
    {
        auto it = blobs.find(name);
        if (it == blobs.end())
            return false;

        auto blob = blobs[name];
        blob->copy_from(data.data(), shape);
        return true;
    }

    bool Graph::feed_input(const std::string& name, const std::shared_ptr<Tensor<float>>& tensor)
    {
        auto it = blobs.find(name);
        if (it == blobs.end())
            return false;

        auto blob = blobs[name];
        blob->copy_from(tensor);
        return true;
    }

    bool Graph::load_output(const std::string& name, std::vector<float>& data)
    {
        auto it = blobs.find(name);
        if (it == blobs.end())
            return false;

        auto blob = blobs[name];
        blob->copy_to(data);
        return true;
    }

    void Graph::forward()
    {
        context->begin();
        for (std::string& layer : sorted_layers)
        {
            auto op = operators[layer];
            std::vector<std::shared_ptr<Blob>> input_blobs;
            std::vector<std::shared_ptr<Blob>> output_blobs;

            for (auto& name : op->get_input_names())
            {
                input_blobs.push_back(blobs[name]);
            }
            for (auto& name : op->get_output_names())
            {
                output_blobs.push_back(blobs[name]);
            }
            op->forward(input_blobs, output_blobs);
        }
        context->end();
        context->submit();
    }

#if ALLOW_PARTIAL_FORWARD
    std::string Graph::partial_forward(uint32_t op_index)
    {
        context->begin();

        uint32_t layer_index = 0;
        std::string output_name = "";
        for (std::string& layer : sorted_layers)
        {
            layer_index += 1;
            if (layer_index > op_index)
                break;

            std::cout << layer_index << ": " << layer << std::endl;
            auto op = operators[layer];
            std::vector<std::shared_ptr<Blob>> input_blobs;
            std::vector<std::shared_ptr<Blob>> output_blobs;

            for (auto& name : op->get_input_names())
            {
                input_blobs.push_back(blobs[name]);
            }
            for (auto& name : op->get_output_names())
            {
                output_blobs.push_back(blobs[name]);
            }
            op->forward(input_blobs, output_blobs);

            output_name = op->get_output_names()[0];
        }
        context->end();
        context->submit();
        return output_name;
    }
#endif
}
